package first.group.ssh.keyexchange;

import first.group.ssh.CipherSuite;
import first.group.ssh.ConnectionCallback;
import first.group.ssh.PacketConnectStage;
import first.group.ssh.SSHConsts;
import first.group.ssh.encrypt.Encryption;
import first.group.ssh.hostkey.HostKeyAlgorithm;
import first.group.ssh.mac.MacAlgorithm;
import first.group.ssh.packet.KeyExchangeInitPacket;
import first.group.ssh.packet.Packet;
import java.io.IOException;
import java.util.NoSuchElementException;
import java.util.Random;
import java.util.logging.Level;
import java.util.logging.Logger;

public class KeyAlgoNegotiator extends PacketConnectStage {

    private static final Logger logger = Logger.getLogger("KeyAlgoNegotiator");

    private final byte[] cookie;
    private final ConnectionCallback callback;
    private CipherSuite suite;
    private final KeyExchangeInitPacket clientInitPacket;
    private KeyExchangeInitPacket serverInitPacket;

    public KeyAlgoNegotiator(Random random, ConnectionCallback callback) {
        this.callback = callback;
        cookie = new byte[16];
        random.nextBytes(cookie);
        clientInitPacket = new KeyExchangeInitPacket(cookie);
    }

    @Override
    public void init() {
        callback.appendPacket(clientInitPacket);
    }

    @Override
    public boolean shouldRead() {
        return true;
    }

    public CipherSuite getSuite() {
        return suite;
    }

    public KeyExchangeInitPacket getClientInitPacket() {
        return clientInitPacket;
    }

    public KeyExchangeInitPacket getServerInitPacket() {
        return serverInitPacket;
    }

    private void findBestAlgorithms(KeyExchangeInitPacket packet) throws IOException {
        if (packet.firstKexPacketFollows) {
            throw new UnsupportedOperationException("Not supported.");
        }
        try {
            var keyExchangeMethod = KeyExchangeMethod.methods
                    .stream()
                    .filter(o -> packet.keyExchangeMethods.contains(o.getName()))
                    .findFirst()
                    .get();
            var hostKeyAlgorithm = HostKeyAlgorithm.algos
                    .stream()
                    .filter(o -> packet.hostKeyAlgorithms.contains(o.getName()))
                    .findFirst()
                    .get();
            var clientToServerEncryption = Encryption.encryptions
                    .stream()
                    .filter(o -> packet.clientToServerEncryptions.contains(o.getName()))
                    .findFirst()
                    .get();
            var serverToClientEncryption = Encryption.encryptions
                    .stream()
                    .filter(o -> packet.serverToClientEncryptions.contains(o.getName()))
                    .findFirst()
                    .get();
            var clientToServerMac = MacAlgorithm.algos
                    .stream()
                    .filter(o -> packet.clientToServerMacs.contains(o.getName()))
                    .findFirst()
                    .get();
            var serverToClientMac = MacAlgorithm.algos
                    .stream()
                    .filter(o -> packet.serverToClientMacs.contains(o.getName()))
                    .findFirst()
                    .get();
            suite = new CipherSuite(
                    keyExchangeMethod,
                    hostKeyAlgorithm,
                    clientToServerEncryption,
                    serverToClientEncryption,
                    clientToServerMac,
                    serverToClientMac
            );
        } catch (NoSuchElementException ex) {
            throw new IOException("No cipher suite found.");
        }
        logger.log(Level.FINE, "Selected suite: {0}", suite);
    }

    @Override
    public void read(Packet packet) throws IOException {
        if (packet.getType() != SSHConsts.SSH_MSG_KEXINIT) {
            throw new IOException("Want packet type " + SSHConsts.SSH_MSG_KEXINIT + ", but get" + packet.getType());
        }
        var initPacket = (KeyExchangeInitPacket) packet;
        serverInitPacket = initPacket;
        findBestAlgorithms(serverInitPacket);
        callback.finish();
    }
}
